#include "cgroup_manager.h"
#include <spdlog/spdlog.h>
#include <fstream>
#include <filesystem>
#include <thread>
#include <signal.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <dirent.h>
#include <errno.h>
#include <unistd.h>
#include <sstream>
#include <chrono>
#include <stdexcept>
#include <system_error>
#include <atomic>

namespace fs = std::filesystem;

namespace judge {

namespace {
    // 读取cgroup文件内容
    std::string readFile(const std::string& path) {
        std::ifstream ifs(path);
        if (!ifs) {
            throw std::runtime_error("Failed to open file: " + path + " - " + strerror(errno));
        }
        std::string content;
        std::getline(ifs, content);
        return content;
    }

    // 写入cgroup文件
    void writeFile(const std::string& path, const std::string& value) {
        std::ofstream ofs(path);
        if (!ofs) {
            throw std::runtime_error("Failed to open file: " + path + " - " + strerror(errno));
        }
        ofs << value;
        if (!ofs) {
            throw std::runtime_error("Failed to write to file: " + path + " - " + strerror(errno));
        }
        ofs.flush();
    }

    // 检查文件是否存在且可写
    bool isWritable(const std::string& path) {
        return access(path.c_str(), W_OK) == 0;
    }

    // 等待文件变为可写
    bool waitForWritable(const std::string& path, int timeout_ms = 1000) {
        for (int i = 0; i < timeout_ms / 10; i++) {
            if (isWritable(path)) {
                return true;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
        return false;
    }

    // 检查cgroup v2是否已挂载
    bool isCgroupV2Mounted() {
        std::ifstream mounts("/proc/mounts");
        std::string line;
        while (std::getline(mounts, line)) {
            if (line.find("cgroup2") != std::string::npos && 
                line.find("/sys/fs/cgroup") != std::string::npos) {
                return true;
            }
        }
        return false;
    }

    // 等待cgroup v2挂载
    bool waitForCgroupV2(int timeout_ms = 5000) {
        for (int i = 0; i < timeout_ms / 100; i++) {
            if (isCgroupV2Mounted()) {
                return true;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
        return false;
    }
}

static std::atomic<int> g_cgroup_instance_counter{0};

CgroupManager::CgroupManager() {
    spdlog::set_level(spdlog::level::err);

    try {
        // 等待cgroup v2挂载
        if (!waitForCgroupV2()) {
            throw std::runtime_error("Cgroup v2 is not mounted");
        }

        // 初始化cgroup路径
        cgroup_path_ = "/sys/fs/cgroup/sandbox/sandbox_" + std::to_string(getpid()) + "_" + std::to_string(++g_cgroup_instance_counter);
        spdlog::debug("Initializing cgroup at path: {}", cgroup_path_);

        // 创建父目录
        if (mkdir("/sys/fs/cgroup/sandbox", 0755) == -1 && errno != EEXIST) {
            throw std::runtime_error("Failed to create sandbox cgroup directory: " + std::string(strerror(errno)));
        }

        // 等待父目录变为可写
        if (!waitForWritable("/sys/fs/cgroup/sandbox")) {
            throw std::runtime_error("Parent cgroup directory is not writable");
        }

        // 在父cgroup中启用控制器
        try {
            writeFile("/sys/fs/cgroup/cgroup.subtree_control", "+cpu +memory +pids");
            writeFile("/sys/fs/cgroup/sandbox/cgroup.subtree_control", "+cpu +memory +pids");
        } catch (const std::exception& e) {
            spdlog::warn("Failed to enable controllers: {}", e.what());
        }

        // 创建此沙箱实例的cgroup目录
        if (mkdir(cgroup_path_.c_str(), 0755) == -1) {
            if (errno != EEXIST) {
                throw std::runtime_error("Failed to create cgroup directory: " + std::string(strerror(errno)));
            }
            // 如果目录已存在，先清理
            cleanup();
            if (mkdir(cgroup_path_.c_str(), 0755) == -1) {
                throw std::runtime_error("Failed to create cgroup directory after cleanup: " + std::string(strerror(errno)));
            }
        }

        // 等待cgroup目录变为可写
        if (!waitForWritable(cgroup_path_)) {
            throw std::runtime_error("Cgroup directory is not writable");
        }

        // 初始化子系统
        initSubsystems();
    } catch (const std::exception& e) {
        spdlog::warn("Cgroup initialization failed: {}. Running without cgroup limits.", e.what());
        disabled_ = true;
    }
}

CgroupManager::~CgroupManager() {
    cleanup();
}

void CgroupManager::initSubsystems() {
    if (disabled_) return;
    try {
        // 启用所有必要的控制器
        writeFile(cgroup_path_ + "/cgroup.subtree_control", "+cpu +memory +pids");

        // 设置CPU限制
        // cpu.max格式: $MAX_QUOTA $PERIOD
        // 例如: "100000 1000000" 表示在1秒内最多使用100ms的CPU时间
        if (cpu_time_limit_.count() > 0) {
            int period_us = 100000;  // 100ms周期，提供更精确的控制
            int quota_us = cpu_time_limit_.count() * 1000;  // 转换为微秒
            writeFile(cgroup_path_ + "/cpu.max", 
                     std::to_string(quota_us) + " " + std::to_string(period_us));
        }

        // 设置内存限制
        if (memory_limit_ > 0) {
            // 设置内存限制，包括缓存
            writeFile(cgroup_path_ + "/memory.max", std::to_string(memory_limit_));
            // 禁用swap
            writeFile(cgroup_path_ + "/memory.swap.max", "0");
            // 设置OOM优先级
            writeFile(cgroup_path_ + "/memory.oom.group", "1");
            // 设置内存使用优先级
            writeFile(cgroup_path_ + "/memory.low", "0");
            writeFile(cgroup_path_ + "/memory.high", std::to_string(memory_limit_ * 95 / 100));  // 95%触发软限制
        }

        // 设置进程数限制
        writeFile(cgroup_path_ + "/pids.max", "32");  // 限制更严格

        spdlog::debug("Cgroup subsystems initialized successfully");
    } catch (const std::exception& e) {
        throw std::runtime_error("Failed to initialize cgroup subsystems: " + std::string(e.what()));
    }
}

void CgroupManager::addProcess(pid_t pid) {
    if (disabled_) return;
    try {
        // 等待cgroup.procs文件变为可写
        std::string procs_file = cgroup_path_ + "/cgroup.procs";
        if (!waitForWritable(procs_file)) {
            throw std::runtime_error("Cgroup procs file is not writable");
        }

        // 将进程添加到cgroup
        writeFile(procs_file, std::to_string(pid));
        
        // 验证进程是否成功添加到cgroup
        bool success = false;
        for (int i = 0; i < 50; i++) {  // 最多等待500ms
            try {
                std::string content = readFile(procs_file);
                std::istringstream iss(content);
                std::string line;
                while (std::getline(iss, line)) {
                    if (line == std::to_string(pid)) {
                        success = true;
                        break;
                    }
                }
                if (success) break;
            } catch (...) {}
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }

        if (!success) {
            throw std::runtime_error("Failed to verify process addition to cgroup");
        }

        // 确保资源限制生效
        std::string memory_current = readFile(cgroup_path_ + "/memory.current");
        std::string cpu_stat = readFile(cgroup_path_ + "/cpu.stat");
        
        if (memory_current.empty() || cpu_stat.empty()) {
            throw std::runtime_error("Failed to read resource usage stats");
        }

        spdlog::debug("Process {} successfully added to cgroup", pid);
    } catch (const std::exception& e) {
        throw std::runtime_error("Failed to add process to cgroup: " + std::string(e.what()));
    }
}

std::pair<std::chrono::microseconds, std::size_t> CgroupManager::getResourceUsage() {
    if (disabled_) return {std::chrono::microseconds(0), 0};
    try {
        // 读取CPU使用时间
        std::string cpu_stat = readFile(cgroup_path_ + "/cpu.stat");
        std::istringstream iss(cpu_stat);
        std::string line;
        std::chrono::microseconds cpu_usage(0);
        while (std::getline(iss, line)) {
            if (line.find("usage_usec") != std::string::npos) {
                std::size_t pos = line.find_last_of(" \t");
                if (pos != std::string::npos) {
                    cpu_usage = std::chrono::microseconds(std::stoull(line.substr(pos + 1)));
                }
                break;
            }
        }

        // 读取内存使用量
        std::string memory_current = readFile(cgroup_path_ + "/memory.current");

        // 读取内存峰值
        std::string memory_peak = readFile(cgroup_path_ + "/memory.peak");
        std::size_t peak_usage = std::stoull(memory_peak);

        // 使用峰值作为实际内存使用量
        return {cpu_usage, peak_usage};
    } catch (const std::exception& e) {
        spdlog::error("Failed to get resource usage: {}", e.what());
        return {std::chrono::microseconds(0), 0};
    }
}

void CgroupManager::cleanup() {
    if (disabled_) return;
    try {
        // 终止所有进程
        std::string procs_file = cgroup_path_ + "/cgroup.procs";
        if (fs::exists(procs_file)) {
            std::string content = readFile(procs_file);
            std::istringstream iss(content);
            std::string line;
            while (std::getline(iss, line)) {
                try {
                    pid_t pid = std::stoi(line);
                    kill(pid, SIGKILL);
                } catch (...) {}
            }
        }

        // 等待所有进程退出
        for (int i = 0; i < 50; i++) {  // 最多等待500ms
            try {
                std::string content = readFile(procs_file);
                if (content.empty()) break;
            } catch (...) {
                break;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }

        // 重置所有控制器
        try {
            writeFile(cgroup_path_ + "/cgroup.kill", "1");
        } catch (...) {}

        // 删除cgroup目录
        if (fs::exists(cgroup_path_)) {
            std::error_code ec;
            fs::remove_all(cgroup_path_, ec); // ignore permission errors
        }
    } catch (const std::exception& e) {
        spdlog::error("Failed to cleanup cgroup: {}", e.what());
    }
}

void CgroupManager::writeToFile(const std::string& path, const std::string& content) {
    std::ofstream file(path);
    if (!file) {
        throw std::runtime_error("Failed to open file: " + path);
    }
    file << content;
    if (!file) {
        throw std::runtime_error("Failed to write to file: " + path);
    }
}

void CgroupManager::setCPUTimeLimit(std::chrono::milliseconds ms) {
    cpu_time_limit_ = ms;
    if (disabled_) return;
    try {
        int period_us = 1000000;  // 1秒
        int quota_us = ms.count() * 1000;  // 转换为微秒
        writeFile(cgroup_path_ + "/cpu.max", 
                 std::to_string(quota_us) + " " + std::to_string(period_us));
    } catch (const std::exception& e) {
        spdlog::warn("Failed to set CPU time limit: {}", e.what());
    }
}

void CgroupManager::setMemoryLimit(std::size_t bytes) {
    memory_limit_ = bytes;
    if (disabled_) return;
    try {
        // 设置内存限制，包括缓存
        writeFile(cgroup_path_ + "/memory.max", std::to_string(bytes));
        // 禁用swap
        writeFile(cgroup_path_ + "/memory.swap.max", "0");
        // 设置内存使用优先级
        writeFile(cgroup_path_ + "/memory.low", "0");
        writeFile(cgroup_path_ + "/memory.high", std::to_string(bytes));
    } catch (const std::exception& e) {
        spdlog::warn("Failed to set memory limit: {}", e.what());
    }
}

} // namespace judge 